import re
import socket
from urllib.parse import urlparse

from ultralytics.cfg import get_save_dir

from models import common
from models.Detect_Blink_Belt_Thread import BlinkBeltThread
from models.Yolo_Classify_Thread import ClassifyThread
from models.Yolo_Detect_Thread import DetectThread
from models.Yolo_Pose_Thread import PoseThread
from models.Yolo_Track_Thread import TrackThread
from ui.VsWindow import Ui_mainWindow
from ui.rtspDialog import CustomMessageBox
from utils import glo, get_cfg, DEFAULT_CFG, get_file_path, update_yaml, DEFAULT_CFG_PATH, get_folder_path, \
    get_model_path

import json
import os
import shutil
import cv2
import numpy as np
import torch
import importlib
from PySide6.QtGui import QPixmap, QImage, QMouseEvent, QGuiApplication, QColor
from PySide6.QtWidgets import QFileDialog, QGraphicsDropShadowEffect, QFrame, QPushButton
from PySide6.QtUiTools import loadUiType
from PySide6.QtCore import QTimer, Qt, QPropertyAnimation, QEasingCurve, \
    QParallelAnimationGroup, QPoint, Signal
from PySide6 import QtCore, QtGui
from qfluentwidgets import RoundMenu, MenuAnimationType, Action

from ui.webCamera import Camera, WebcamThread
from utils.custom_grips import CustomGrip

glo._init()
glo.set_value('model_name', "detect_blink_belt yolo_detect yolo_classify yolo_track yolo_pose")
GLOBAL_WINDOW_STATE = True
PATH_YOLO_SHOW = os.path.join("ui/VSWINDOW.ui")
formType, baseType = loadUiType(PATH_YOLO_SHOW)
WIDTH_LEFT_BOX_STANDARD = 160
WIDTH_LEFT_BOX_EXTENDED = 160
WIDTH_LOGO = 120
KEYS_LEFT_BOX_MENU = ['src_menu', 'src_setting', 'src_webcam', 'src_folder', 'src_camera', 'src_vsmode', 'src_setting']


# 窗口类 动态加载UI文件 和 Ui_mainWindow
class VsShow(formType, baseType, Ui_mainWindow):
    def __init__(self):
        super().__init__()
        self.current_path = os.getcwd()
        self.inputPath = None
        self.default_cfg = get_cfg(DEFAULT_CFG)
        self.allModelNames = ["detect_blink_belt", "yolo_detect", "yolo_classify", "yolo_track", "yolo_pose"]
        # --- 加载UI --- #
        self.setupUi(self)
        self.setAttribute(Qt.WA_TranslucentBackground, True)  # 透明背景
        self.setWindowFlags(Qt.FramelessWindowHint)  # 无头窗口
        # --- 加载UI --- #

        # 初始化侧边栏
        self.initSiderWidget()

        # --- 最大化 最小化 关闭 --- #
        self.maximizeButton.clicked.connect(self.maxorRestore)
        self.minimizeButton.clicked.connect(self.showMinimized)
        self.closeButton.clicked.connect(self.close)
        self.topbox.doubleClickFrame.connect(self.maxorRestore)
        # --- 最大化 最小化 关闭 --- #

        # --- 播放 暂停 停止 --- #
        self.playIcon = QtGui.QIcon()
        self.playIcon.addPixmap(QtGui.QPixmap(f"{self.current_path}/images/newsize/play.png"), QtGui.QIcon.Normal,
                                QtGui.QIcon.Off)
        self.playIcon.addPixmap(QtGui.QPixmap(f"{self.current_path}/images/newsize/pause.png"), QtGui.QIcon.Active,
                                QtGui.QIcon.On)
        self.playIcon.addPixmap(QtGui.QPixmap(f"{self.current_path}/images/newsize/pause.png"),
                                QtGui.QIcon.Selected, QtGui.QIcon.On)
        self.run_button.setCheckable(True)
        self.run_button.setIcon(self.playIcon)
        # --- 播放 暂停 停止 --- #

        # --- 侧边栏缩放 --- #
        self.src_menu.clicked.connect(self.scaleMenu)  # hide menu button
        self.src_setting.clicked.connect(self.scalSetting)  # setting button
        # --- 侧边栏缩放 --- #

        # --- 自动加载/动态改变 PT 模型 --- #
        self.pt_Path = f"{self.current_path}/weights/"
        self.pt_list = os.listdir(f'{self.current_path}/weights/')
        self.pt_list = [file for file in self.pt_list if file.endswith('.pt')]
        self.pt_list.sort(key=lambda x: os.path.getsize(f'{self.current_path}/weights/' + x))
        self.model_box1.clear()
        self.model_box1.addItems([file.split('.')[0] for file in self.pt_list])
        self.model_box2.clear()
        self.model_box2.addItems([file.split('.')[0] for file in self.pt_list])
        self.qtimer_search = QTimer(self)
        self.qtimer_search.timeout.connect(lambda: self.loadModels())
        self.qtimer_search.start(2000)
        self.model_box1.currentTextChanged.connect(lambda: self.changeModel("left"))
        self.model_box2.currentTextChanged.connect(lambda: self.changeModel("right"))
        # --- 自动加载/动态改变 PT 模型 --- #

        # --- 超参数调整 --- #
        self.iou_spinbox.valueChanged.connect(lambda x: self.changeValue(x, 'iou_spinbox'))  # iou box
        self.iou_slider.valueChanged.connect(lambda x: self.changeValue(x, 'iou_slider'))  # iou scroll bar
        self.conf_spinbox.valueChanged.connect(lambda x: self.changeValue(x, 'conf_spinbox'))  # conf box
        self.conf_slider.valueChanged.connect(lambda x: self.changeValue(x, 'conf_slider'))  # conf scroll bar
        self.speed_spinbox.valueChanged.connect(lambda x: self.changeValue(x, 'speed_spinbox'))  # speed box
        self.speed_slider.valueChanged.connect(lambda x: self.changeValue(x, 'speed_slider'))  # speed scroll bar
        self.line_spinbox.valueChanged.connect(lambda x: self.changeValue(x, 'line_spinbox'))  # line box
        self.line_slider.valueChanged.connect(lambda x: self.changeValue(x, 'line_slider'))  # line slider
        # --- 超参数调整 --- #

        # --- 导入 图片/视频、调用摄像头、导入文件夹（批量处理）、调用网络摄像头 --- #
        self.src_img.clicked.connect(self.selectFile)
        # 对比模型模式 不支持同时读取摄像头流
        self.src_webcam.clicked.connect(self.selectWebcam)
        self.src_folder.clicked.connect(self.selectFolder)
        self.src_camera.clicked.connect(self.selectRtsp)
        # --- 导入 图片/视频、调用摄像头、导入文件夹（批量处理）、调用网络摄像头 --- #

        # --- 导入模型、 导出结果 --- #
        self.import_button.clicked.connect(self.importModel)
        self.save_status_button.clicked.connect(self.saveStatus)
        self.save_button.clicked.connect(self.saveResult)
        self.save_button.setEnabled(False)
        # --- 导入模型、 导出结果 --- #

        # --- 视频、图片 预览 --- #
        self.main_leftbox.setAlignment(QtCore.Qt.AlignCenter | QtCore.Qt.AlignVCenter)
        self.main_rightbox.setAlignment(QtCore.Qt.AlignCenter | QtCore.Qt.AlignVCenter)
        # --- 视频、图片 预览 --- #

        # --- 状态栏 初始化 --- #
        # 状态栏阴影效果
        self.shadowStyle(self.mainBody, QColor(0, 0, 0, 38), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Class_QF1, QColor(142, 197, 252), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Target_QF1, QColor(159, 172, 230), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Fps_QF1, QColor(170, 128, 213), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Model_QF1, QColor(162, 129, 247), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Class_QF2, QColor(142, 197, 252), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Target_QF2, QColor(159, 172, 230), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Fps_QF2, QColor(170, 128, 213), top_bottom=['top', 'bottom'])
        self.shadowStyle(self.Model_QF2, QColor(162, 129, 247), top_bottom=['top', 'bottom'])

        # 左侧状态栏默认显示
        self.model_name1 = self.model_box1.currentText()  # 获取默认 model
        self.Class_num1.setText('--')
        self.Target_num1.setText('--')
        self.fps_label1.setText('--')
        self.Model_label1.setText(self.model_name1)

        # 右侧状态栏默认显示
        self.model_name2 = self.model_box2.currentText()  # 获取默认 model
        self.Class_num2.setText('--')
        self.Target_num2.setText('--')
        self.fps_label2.setText('--')
        self.Model_label2.setText(self.model_name2)

        self.left_Threads, self.right_Threads = {}, {}

        # detect_blink_belt QThread 初始化
        self.left_Threads["detect_blink_belt"] = BlinkBeltThread()
        self.right_Threads["detect_blink_belt"] = BlinkBeltThread()
        self.initModel("detect_blink_belt")

        # --- yolo_detect QThread --- #
        self.left_Threads["yolo_detect"] = DetectThread()
        self.right_Threads["yolo_detect"] = DetectThread()
        self.initModel("yolo_detect")

        # --- yolo_classify QThread --- #
        self.left_Threads["yolo_classify"] = ClassifyThread()
        self.right_Threads["yolo_classify"] = ClassifyThread()
        self.initModel("yolo_classify")

        # --- yolo_track QThread --- #
        self.left_Threads["yolo_track"] = TrackThread()
        self.right_Threads["yolo_track"] = TrackThread()
        self.initModel("yolo_track")

        # --- yolo_pose QThread --- #
        self.left_Threads["yolo_pose"] = PoseThread()
        self.right_Threads["yolo_pose"] = PoseThread()
        self.initModel("yolo_pose")

        # --- 状态栏 初始化 --- #
        self.run_button.clicked.connect(self.runorContinue)
        # --- 开始 / 停止 --- #

        # --- Setting栏 初始化 --- #
        self.loadConfig()
        # --- Setting栏 初始化 --- #

        # --- MessageBar Init --- #
        self.showStatus("Welcome to YOLOSHOW")
        # --- MessageBar Init --- #

    def initSiderWidget(self):
        # --- 侧边栏 --- #
        self.leftBox.setFixedWidth(WIDTH_LEFT_BOX_STANDARD)
        # logo
        self.logo.setFixedSize(WIDTH_LOGO, WIDTH_LOGO)

        # 将左侧菜单栏的按钮固定宽度
        for child_left_box_widget in self.leftbox_bottom.children():

            if isinstance(child_left_box_widget, QFrame):
                child_left_box_widget.setFixedWidth(WIDTH_LEFT_BOX_EXTENDED)

                for child_left_box_widget_btn in child_left_box_widget.children():
                    if isinstance(child_left_box_widget_btn, QPushButton):
                        child_left_box_widget_btn.setFixedWidth(WIDTH_LEFT_BOX_EXTENDED)

    def loadModelConfig(self, thread, mode="left"):
        if mode == "left":
            thread.new_model_name = f'{self.current_path}/weights/' + self.model_box1.currentText() + '.pt'
            thread.send_output.connect(lambda x: self.showImg(x, self.main_leftbox, 'img'))
            thread.send_msg.connect(lambda x: self.showStatus(x))
            thread.send_fps.connect(lambda x: self.fps_label1.setText(str(x)))
            thread.send_class_num.connect(lambda x: self.Class_num1.setText(str(x)))
            thread.send_target_num.connect(lambda x: self.Target_num1.setText(str(x)))
        else:
            thread.new_model_name = f'{self.current_path}/weights/' + self.model_box2.currentText() + '.pt'
            thread.send_output.connect(lambda x: self.showImg(x, self.main_rightbox, 'img'))
            thread.send_progress.connect(lambda x: self.progress_bar.setValue(x))
            thread.send_fps.connect(lambda x: self.fps_label2.setText(str(x)))
            thread.send_class_num.connect(lambda x: self.Class_num2.setText(str(x)))
            thread.send_target_num.connect(lambda x: self.Target_num2.setText(str(x)))
        thread.max_progress_value = self.progress_bar.maximum()

    # 初始化模型
    def initModel(self, model_name=None, mode="all"):
        if model_name == "detect_blink_belt":
            if mode == "all" or mode == "left":
                # 左侧模型加载 第一个模型来控制消息
                self.loadModelConfig(self.left_Threads["detect_blink_belt"], mode="left")
            if mode == "all" or mode == "right":
                # 右侧模型加载
                self.loadModelConfig(self.right_Threads["detect_blink_belt"], mode="right")
        elif model_name == "yolo_detect":
            if mode == "all" or mode == "left":
                self.loadModelConfig(self.left_Threads["yolo_detect"], mode="left")
            if mode == "all" or mode == "right":
                self.loadModelConfig(self.right_Threads["yolo_detect"], mode="right")
        elif model_name == "yolo_classify":
            if mode == "all" or mode == "left":
                self.loadModelConfig(self.left_Threads["yolo_classify"], mode="left")
            if mode == "all" or mode == "right":
                self.loadModelConfig(self.right_Threads["yolo_classify"], mode="right")
        elif model_name == "yolo_track":
            if mode == "all" or mode == "left":
                self.loadModelConfig(self.left_Threads["yolo_track"], mode="left")
            if mode == "all" or mode == "right":
                self.loadModelConfig(self.right_Threads["yolo_track"], mode="right")
        # --- YOLOv5-seg QThread --- #
        elif model_name == "yolo_pose":
            if mode == "all" or mode == "left":
                self.loadModelConfig(self.left_Threads["yolo_pose"], mode="left")
            if mode == "all" or mode == "right":
                self.loadModelConfig(self.right_Threads["yolo_pose"], mode="right")

    # 阴影效果
    def shadowStyle(self, widget, Color, top_bottom=None):
        shadow = QGraphicsDropShadowEffect(self)
        if 'top' in top_bottom and 'bottom' not in top_bottom:
            shadow.setOffset(0, -5)
        elif 'bottom' in top_bottom and 'top' not in top_bottom:
            shadow.setOffset(0, 5)
        else:
            shadow.setOffset(5, 5)
        shadow.setBlurRadius(10)  # 阴影半径
        shadow.setColor(Color)  # 阴影颜色
        widget.setGraphicsEffect(shadow)

    # 侧边栏缩放
    def scaleMenu(self):
        standard = 160
        maxExtend = 160
        width = self.leftBox.width()
        if width == 120:
            widthExtended = maxExtend
        else:
            widthExtended = standard
        # animation
        self.animation = QPropertyAnimation(self.leftBox, b"minimumWidth")
        self.animation.setDuration(500)  # ms
        self.animation.setStartValue(width)
        self.animation.setEndValue(widthExtended)
        self.animation.setEasingCurve(QEasingCurve.InOutQuint)
        self.animation.start()

    # 设置栏缩放
    def scalSetting(self):
        # GET WIDTH
        widthSettingBox = self.settingBox.width()  # right set column width
        widthLeftBox = self.leftBox.width()  # left column length
        maxExtend = 220
        standard = 0
        # SET MAX WIDTH
        if widthSettingBox == 0:
            widthExtended = maxExtend
            self.mainbox.setStyleSheet("""
                                  QFrame#mainbox{
                                    border: 1px solid rgba(0, 0, 0, 15%);
                                    border-bottom-left-radius: 0;
                                    border-bottom-right-radius: 0;
                                    border-radius:30%;
                                    background-color: qlineargradient(x1:0, y1:0, x2:1 , y2:0, stop:0 white, stop:0.9 #8EC5FC, stop:1 #E0C3FC);
                                }
                              """)
        else:
            widthExtended = standard
            self.mainbox.setStyleSheet("""
                                  QFrame#mainbox{
                                    border: 1px solid rgba(0, 0, 0, 15%);
                                    border-bottom-left-radius: 0;
                                    border-bottom-right-radius: 0;
                                    border-radius:30%;
                                }
                              """)

        # ANIMATION LEFT BOX
        self.left_box = QPropertyAnimation(self.leftBox, b"minimumWidth")
        self.left_box.setDuration(100)
        self.left_box.setStartValue(widthLeftBox)
        self.left_box.setEndValue(68)
        self.left_box.setEasingCurve(QEasingCurve.InOutQuart)

        # ANIMATION SETTING BOX
        self.setting_box = QPropertyAnimation(self.settingBox, b"minimumWidth")
        self.setting_box.setDuration(100)
        self.setting_box.setStartValue(widthSettingBox)
        self.setting_box.setEndValue(widthExtended)
        self.setting_box.setEasingCurve(QEasingCurve.InOutQuart)

        # SET QSS Change
        self.qss_animation = QPropertyAnimation(self.mainbox, b"styleSheet")
        self.qss_animation.setDuration(80)
        self.qss_animation.setStartValue("""
            QFrame#mainbox {
                border: 1px solid rgba(0, 0, 0, 15%);
                border-bottom-left-radius: 0;
                border-bottom-right-radius: 0;
                border-radius:30%;
                background-color: qlineargradient(x1:0, y1:0, x2:1, y2:0, stop:0 white, stop:0.9 #8EC5FC, stop:1 #E0C3FC);
            }
        """)
        self.qss_animation.setEndValue("""
             QFrame#mainbox {
                border: 1px solid rgba(0, 0, 0, 15%);
                border-bottom-left-radius: 0;
                border-bottom-right-radius: 0;
                border-radius:30%;
            }
        """)
        self.qss_animation.setEasingCurve(QEasingCurve.InOutQuart)

        # GROUP ANIMATION
        self.group = QParallelAnimationGroup()
        self.group.addAnimation(self.left_box)
        self.group.addAnimation(self.setting_box)
        self.group.start()

    # 最大化最小化窗口
    def maxorRestore(self):
        global GLOBAL_WINDOW_STATE
        status = GLOBAL_WINDOW_STATE
        if status:
            self.showMaximized()
            self.maximizeButton.setStyleSheet("""
                          QPushButton:hover{
                               background-color:rgb(139, 29, 31);
                               border-image: url(:/leftbox/images/newsize/scalling.png);
                           }
                      """)
            GLOBAL_WINDOW_STATE = False
        else:
            self.showNormal()
            self.maximizeButton.setStyleSheet("""
                                      QPushButton:hover{
                                           background-color:rgb(139, 29, 31);
                                           border-image: url(:/leftbox/images/newsize/max.png);
                                       }
                                  """)
            GLOBAL_WINDOW_STATE = True

    # 选择照片/视频 并展示
    def selectFile(self):
        # 获取上次选择文件的路径
        file_path = get_file_path(self.default_cfg)

        file, _ = QFileDialog.getOpenFileName(
            self,  # 父窗口对象
            "Select your Image / Video",  # 标题
            file_path,  # 默认打开路径为当前路径
            "Image / Video type (*.jpg *.jpeg *.png *.bmp *.dib *.jpe *.jp2 *.mp4)"  # 选择类型过滤项，过滤内容在括号中
        )
        if file:
            self.inputPath = file
            glo.set_value('inputPath', self.inputPath)
            self.update_model_data(str(self.inputPath))
            # 如果是视频， 显示第一帧
            if ".avi" in self.inputPath or ".mp4" in self.inputPath:
                # 显示第一帧
                self.cap = cv2.VideoCapture(self.inputPath)
                ret, frame = self.cap.read()
                if ret:
                    # rgbImage = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    self.showImg(frame, self.main_leftbox, 'img')
            # 如果是图片 正常显示
            else:
                self.showImg(self.inputPath, self.main_leftbox, 'path')
            self.showStatus('Loaded File：{}'.format(os.path.basename(self.inputPath)))
            update_yaml(DEFAULT_CFG_PATH, "file_path", self.inputPath)  # 更新yaml文件
            self.default_cfg = get_cfg(DEFAULT_CFG)  # 更新cfg参数

    def update_model_data(self, data_path):
        # 更新模型中的data
        if self.left_Threads[self.model_name1] is not None and self.left_Threads[self.model_name1].isRunning:
            self.left_Threads[self.model_name1].update_dataset(data_path)
        if self.right_Threads[self.model_name2] is not None and self.right_Threads[self.model_name2].isRunning:
            self.right_Threads[self.model_name2].update_dataset(data_path)
        self.main_rightbox.clear()
        self.main_leftbox.clear()
        # 开始检测 第一次设置为false
        self.run_button.setChecked(False)

    # 选择摄像头
    def selectWebcam(self):
        try:
            # get the number of local cameras
            cam_num, cams = Camera().get_cam_num()
            if cam_num > 0:
                popMenu = RoundMenu(parent=self)
                popMenu.setFixedWidth(self.leftbox_bottom.width())
                actions = []

                for cam in cams:
                    cam_name = f'Camera_{cam}'
                    actions.append(Action(cam_name))
                    popMenu.addAction(actions[-1])
                    actions[-1].triggered.connect(lambda: self.actionWebcam(cam))
                x = self.webcamBox.mapToGlobal(self.webcamBox.pos()).x()
                y = self.webcamBox.mapToGlobal(self.webcamBox.pos()).y()
                y = y - self.webcamBox.frameGeometry().height() * 2
                pos = QPoint(x, y)
                popMenu.exec(pos, aniType=MenuAnimationType.DROP_DOWN)
            else:
                self.showStatus('No camera found !!!')
        except Exception as e:
            self.showStatus('%s' % e)

    # 调用网络摄像头
    def actionWebcam(self, cam):
        self.showStatus(f'Loading camera：Camera_{cam}')
        self.thread = WebcamThread(cam)
        self.thread.changePixmap.connect(lambda x: self.showImg(x, self.main_leftbox, 'img'))
        self.thread.start()
        self.inputPath = int(cam)
        self.update_model_data(str(self.inputPath))
        # 选择文件夹

    def selectFolder(self):
        folder_path = get_folder_path(self.default_cfg)
        if not os.path.exists(folder_path):
            folder_path = os.getcwd()
        FolderPath = QFileDialog.getExistingDirectory(
            self,
            "Select your Folder",
            folder_path  # 起始目录
        )
        if FolderPath:
            FileFormat = [".mp4", ".mkv", ".avi", ".flv", ".jpg", ".png", ".jpeg", ".bmp", ".dib", ".jpe", ".jp2"]
            Foldername = [(FolderPath + "/" + filename) for filename in os.listdir(FolderPath) for jpgname in
                          FileFormat
                          if jpgname in filename]
            self.inputPath = Foldername
            self.showStatus('Loaded Folder：{}'.format(os.path.basename(FolderPath)))
            update_yaml(DEFAULT_CFG_PATH, "folder_path", self.FolderPath)  # 更新yaml文件
            self.default_cfg = get_cfg(DEFAULT_CFG)  # 更新cfg参数

    # 选择网络摄像头 Rtsp
    def selectRtsp(self):
        # rtsp://rtsp-test-server.viomic.com:554/stream
        rtspDialog = CustomMessageBox(self, mode="VS")
        self.rtspUrl = None
        if rtspDialog.exec():
            self.rtspUrl = rtspDialog.urlLineEdit.text()
        if self.rtspUrl:
            parsed_url = urlparse(self.rtspUrl)
            if parsed_url.scheme == 'rtsp':
                if not self.checkRtspUrl(self.rtspUrl):
                    self.showStatus('Rtsp stream is not available')
                    return False
                self.showStatus(f'Loading Rtsp：{self.rtspUrl}')
                self.rtspThread = WebcamThread(self.rtspUrl)
                self.rtspThread.changePixmap.connect(lambda x: self.showImg(x, self.main_leftbox, 'img'))
                self.rtspThread.start()
                self.inputPath = self.rtspUrl
                self.update_model_data(str(self.inputPath))
            else:
                self.showStatus('URL is not an rtsp stream')
                return False

    # 检测网络摄像头 Rtsp 是否连通
    def checkRtspUrl(self, url, timeout=5):
        try:
            # 解析URL获取主机名和端口
            from urllib.parse import urlparse
            parsed_url = urlparse(url)
            hostname = parsed_url.hostname
            port = parsed_url.port or 554  # RTSP默认端口是554

            # 创建socket对象
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            # 尝试连接
            sock.connect((hostname, port))
            # 关闭socket
            sock.close()
            return True
        except Exception as e:
            return False

    # 显示Label图片
    @staticmethod
    def showImg(img, label, flag):
        try:
            if flag == "path":
                img_src = cv2.imdecode(np.fromfile(img, dtype=np.uint8), -1)
            else:
                img_src = img
            ih, iw, _ = img_src.shape
            w = label.geometry().width()
            h = label.geometry().height()
            # keep original aspect ratio
            if iw / w > ih / h:
                scal = w / iw
                nw = w
                nh = int(scal * ih)
                img_src_ = cv2.resize(img_src, (nw, nh))
            else:
                scal = h / ih
                nw = int(scal * iw)
                nh = h
                img_src_ = cv2.resize(img_src, (nw, nh))

            frame = cv2.cvtColor(img_src_, cv2.COLOR_BGR2RGB)
            img = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[2] * frame.shape[1],
                         QImage.Format_RGB888)
            label.setPixmap(QPixmap.fromImage(img))
        except Exception as e:
            print(repr(e))

    # resize 窗口大小
    def resizeGrip(self):
        self.left_grip.setGeometry(0, 10, 10, self.height())
        self.right_grip.setGeometry(self.width() - 10, 10, 10, self.height())
        self.top_grip.setGeometry(0, 0, self.width(), 10)
        self.bottom_grip.setGeometry(0, self.height() - 10, self.width(), 10)

    # 停止所有运行中的模型
    def quitRuniningModel(self):
        try:
            for thread in self.left_Threads.values():
                if thread is not None and thread.isRunning:
                    thread.quit()
            for thread in self.right_Threads.values():
                if thread is not None and thread.isRunning:
                    thread.quit()
            self.inputPath = None
            self.resignOneSideModel(self.model_name1,"left")
            self.resignTwoSideModel(self.model_name2,"right")
        except Exception:
            pass

    # 在MessageBar显示消息
    def showStatus(self, msg):
        self.message_bar.setText(msg)
        if msg == 'Finish Detection':
            self.run_button.setChecked(False)
            self.progress_bar.setValue(0)
            self.quitRuniningModel()
            self.save_status_button.setEnabled(True)


    def importModel(self):
        # 获取上次选择文件的路径
        self.model_path = get_model_path(self.default_cfg)
        file, _ = QFileDialog.getOpenFileName(
            self,  # 父窗口对象
            "Select your YOLO Model",  # 标题
            self.model_path,  # 默认打开路径为当前路径
            "Model File (*.pt)"  # 选择类型过滤项，过滤内容在括号中
        )
        if file:
            fileptPath = os.path.join(self.pt_Path, os.path.basename(file))
            if not os.path.exists(fileptPath):
                shutil.copy(file, self.pt_Path)
                self.showStatus('Loaded Model：{}'.format(os.path.basename(file)))
                # 更新yaml文件
                update_yaml(DEFAULT_CFG_PATH, "model_path", os.path.dirname(file))
                self.default_cfg = get_cfg(DEFAULT_CFG)  # 更新cfg参数
            else:
                self.showStatus('Model already exists')

    # 导出结果状态判断
    def saveStatus(self):
        if self.save_status_button.checkState() == Qt.CheckState.Unchecked:
            self.showStatus('NOTE: Run image results are not saved.')
            for thread in self.left_Threads.values():
                thread.save_res = False
            for thread in self.right_Threads.values():
                thread.save_res = False
            self.save_button.setEnabled(False)
        elif self.save_status_button.checkState() == Qt.CheckState.Checked:
            self.showStatus('NOTE: Run image results will be saved. Only save left model result in VS Mode !!!')
            for thread in self.left_Threads.values():
                thread.save_res = True
            for thread in self.right_Threads.values():
                thread.save_res = True
            self.save_button.setEnabled(True)

    # 验证文件是否已经保存
    def valFileIsSave(self, save_dir):
        try:
            output_dir = os.path.dirname(save_dir)
            if os.path.exists(output_dir):
                for filename in os.listdir(output_dir):
                    source_path = os.path.join(output_dir, filename)
                    destination_path = os.path.join(self.OutputDir, filename)
                    if os.path.isfile(source_path):
                        shutil.copy(source_path, destination_path)
                self.showStatus('Saved Successfully in {}'.format(self.OutputDir))
            else:
                self.showStatus('Please wait for the result to be generated')
        except Exception as err:
            self.showStatus(f"Error occurred while saving the result: {err}")

    def valIsSave(self, save_dir):
        try:
            if os.path.exists(save_dir):
                shutil.copy(save_dir, self.OutputDir)
                self.showStatus('Saved Successfully in {}'.format(self.OutputDir))
            else:
                self.showStatus('Please wait for the result to be generated')
        except Exception as err:
            self.showStatus(f"Error occurred while saving the result: {err}")

    # 导出结果
    def saveResult(self):
        # 默认保存左侧模型的检测结果
        if not self.left_Threads[self.model_name1].res_status:
            self.showStatus("导出失败..")
            return
        save_dir = get_save_dir(self.default_cfg)  # 保存路径
        is_folder = isinstance(self.inputPath, list)

        if is_folder:
            self.OutputDir = QFileDialog.getExistingDirectory(
                self,  # 父窗口对象
                "Save Results in new Folder",  # 标题
                save_dir,  # 起始目录
            )
            for thread in self.left_Threads.values():
                self.valFileIsSave(thread.save_dir)
        else:
            self.OutputDir, _ = QFileDialog.getSaveFileName(
                self,  # 父窗口对象
                "Save Image/Video",  # 标题
                save_dir,  # 起始目录
                "Image/Vide Type (*.jpg *.jpeg *.png *.bmp *.dib  *.jpe  *.jp2 *.mp4)"  # 选择类型过滤项，过滤内容在括号中
            )
            for thread in self.left_Threads.values():
                self.valIsSave(thread.save_dir)
            # 更新yaml文件
        update_yaml(DEFAULT_CFG_PATH, "save_dir", self.OutputDir)
        self.default_cfg = get_cfg(DEFAULT_CFG)  # 更新cfg参数

    def setThreadCfg(self, cfg, value):
        for thread in self.left_Threads.values():
            if cfg == "iou_thres":
                thread.iou_thres = value / 100
            elif cfg == "conf_thres":
                thread.conf_thres = value / 100
            elif cfg == "speed_thres":
                thread.speed_thres = value
            elif cfg == "line_slider":
                thread.line_thickness = value
        for thread in self.right_Threads.values():
            if cfg == "iou_thres":
                thread.iou_thres = value / 100
            elif cfg == "conf_thres":
                thread.conf_thres = value / 100
            elif cfg == "speed_thres":
                thread.speed_thres = value
            elif cfg == "line_slider":
                thread.line_thickness = value

    # 调整超参数
    def changeValue(self, x, flag):
        if flag == 'iou_spinbox':
            self.iou_slider.setValue(int(x * 100))  # The box value changes, changing the slider
        elif flag == 'iou_slider':
            self.iou_spinbox.setValue(x / 100)  # The slider value changes, changing the box
            self.showStatus('IOU Threshold: %s' % str(x / 100))
            self.setThreadCfg("iou_thres", x)
        elif flag == 'conf_spinbox':
            self.conf_slider.setValue(int(x * 100))
        elif flag == 'conf_slider':
            self.conf_spinbox.setValue(x / 100)
            self.showStatus('Conf Threshold: %s' % str(x / 100))
            self.setThreadCfg("conf_thres", x)
        elif flag == 'speed_spinbox':
            self.speed_slider.setValue(x)
        elif flag == 'speed_slider':
            self.speed_spinbox.setValue(x)
            self.showStatus('Delay: %s ms' % str(x))
            self.setThreadCfg("speed_thres", x)
        elif flag == 'line_spinbox':
            self.line_slider.setValue(x)
        elif flag == 'line_slider':
            self.line_spinbox.setValue(x)
            self.showStatus('Line Width: %s' % str(x))
            self.setThreadCfg("line_slider", x)

    # 加载 Setting 栏
    def loadConfig(self):
        config_file = 'cfg/setting.json'
        if not os.path.exists(config_file):
            iou = 0.45
            conf = 0.25
            delay = 10
            line_thickness = 3
            new_config = {"iou": iou,
                          "conf": conf,
                          "delay": delay,
                          "line_thickness": line_thickness,
                          }
            new_json = json.dumps(new_config, ensure_ascii=False, indent=2)
            with open(config_file, 'w', encoding='utf-8') as f:
                f.write(new_json)
        else:
            config = json.load(open(config_file, 'r', encoding='utf-8'))
            if len(config) != 4:
                iou = 0.45
                conf = 0.25
                delay = 10
                line_thickness = 3
            else:
                iou = config['iou']
                conf = config['conf']
                delay = config['delay']
                line_thickness = config['line_thickness']
        self.iou_spinbox.setValue(iou)
        self.iou_slider.setValue(int(iou * 100))
        self.conf_spinbox.setValue(conf)
        self.conf_slider.setValue(int(conf * 100))
        self.speed_spinbox.setValue(delay)
        self.speed_slider.setValue(delay)
        self.line_spinbox.setValue(line_thickness)
        self.line_slider.setValue(line_thickness)

    # 加载 pt 模型到 model_box
    def loadModels(self):
        pt_list = os.listdir(f'{self.current_path}/weights/')
        pt_list = [file for file in pt_list if file.endswith('.pt')]
        pt_list.sort(key=lambda x: os.path.getsize(f'{self.current_path}/weights/' + x))
        if pt_list != self.pt_list:
            self.pt_list = pt_list
            self.model_box1.clear()
            self.model_box1.addItems([file.split('.')[0] for file in self.pt_list])
            self.model_box2.clear()
            self.model_box2.addItems([file.split('.')[0] for file in self.pt_list])

    # 检查模型是否符合命名要求
    def checkModelName(self, modelname):
        for name in self.allModelNames:
            if modelname in name:
                return True
        return False

    # 重新加载模型
    def reloadModel(self):
        importlib.reload(common)

    # 重新加载左侧、右侧模型
    def resignOneSideModel(self, model_name=None, mode=None):
        if mode == "left":
            del self.left_Threads[model_name]
            self.Class_num1.setText('--')
            self.Target_num1.setText('--')
            self.fps_label2.setText('--')
            if model_name == "detect_blink_belt":
                self.left_Threads["detect_blink_belt"] = BlinkBeltThread()
                self.initModel("detect_blink_belt", "left")
            elif model_name == "yolo_detect":
                self.left_Threads["yolo_detect"] = DetectThread()
                self.initModel("yolo_detect", "left")
            elif model_name == "yolo_classify":
                self.left_Threads["yolo_classify"] = ClassifyThread()
                self.initModel("yolo_classify", "left")
            elif model_name == "yolo_track":
                self.left_Threads["yolo_track"] = TrackThread()
                self.initModel("yolo_track", "left")
            elif model_name == "yolo_pose":
                self.left_Threads["yolo_pose"] = PoseThread()
                self.initModel("yolo_pose", "left")
        else:
            del self.right_Threads[model_name]
            self.Class_num2.setText('--')
            self.Target_num2.setText('--')
            self.fps_label2.setText('--')
            if model_name == "detect_blink_belt":
                self.right_Threads["detect_blink_belt"] = BlinkBeltThread()
                self.initModel("detect_blink_belt", "right")
            elif model_name == "yolo_detect":
                self.right_Threads["yolo_detect"] = DetectThread()
                self.initModel("yolo_detect", "right")
            elif model_name == "yolo_classify":
                self.right_Threads["yolo_classify"] = ClassifyThread()
                self.initModel("yolo_classify", "right")
            elif model_name == "yolo_track":
                self.right_Threads["yolo_track"] = TrackThread()
                self.initModel("yolo_track", "right")
            elif model_name == "yolo_pose":
                self.right_Threads["yolo_pose"] = PoseThread()
                self.initModel("yolo_pose", "right")

    # 停止所有模型,重新加载某一侧current_model_name模型 和另一侧模型
    def stopOtherModel(self, mode="all"):
        self.inputPath = None
        if mode == "left" or mode == "all":
            self.stopOneSideModel(self.model_name1, mode="left")
        if mode == "right" or mode == "all":
            self.stopOneSideModel(self.model_name2, mode="right")
        self.save_status_button.setEnabled(True)
        self.progress_bar.setValue(0)
        self.main_leftbox.clear()  # clear image display
        self.main_rightbox.clear()
        self.runModel(True)

    # 停止一侧模型
    def stopOneSideModel(self, current_model_name=None, mode=None):
        all_model_names = self.allModelNames
        if mode == "left":
            for model_name in all_model_names:
                if model_name != current_model_name:
                    if self.left_Threads[model_name] is not None and self.left_Threads[model_name].isRunning():
                        self.left_Threads[model_name].quit_flag = True
                        self.left_Threads[model_name].finished.connect(
                            lambda: self.resignOneSideModel(current_model_name, "left"))
        else:
            for model_name in all_model_names:
                if model_name != current_model_name:
                    if self.right_Threads[model_name] is not None and self.right_Threads[model_name].isRunning():
                        self.right_Threads[model_name].quit_flag = True
                        self.right_Threads[model_name].finished.connect(
                            lambda: self.resignOneSideModel(current_model_name, "right"))

    # Model 变化
    def changeModel(self, mode=None):
        if mode == "left":
            # 左侧模型
            self.model_name1 = self.model_box1.currentText()
            self.Model_label1.setText(self.model_name1)  # 修改状态栏显示
            left_count = 0
            for key in self.left_Threads:
                if self.model_name1 == key:
                    left_count += 1
                    self.left_Threads[key].new_model_name = f'{self.current_path}/weights/' + key + '.pt'
                    glo.set_value('model_name1', key)
                    self.reloadModel()
            if left_count == 0:
                self.showStatus('The current left-model is not supported')
            else:
                self.stopOtherModel()
        else:
            # 右侧模型
            self.model_name2 = self.model_box2.currentText()
            self.Model_label2.setText(self.model_name2)
            right_count = 0
            for key in self.right_Threads:
                if self.model_name2 == key:
                    right_count += 1
                    self.right_Threads[key].new_model_name = f'{self.current_path}/weights/' + key + '.pt'
                    glo.set_value('model_name2', key)
                    self.reloadModel()
            if right_count == 0:
                self.showStatus('The current right-model is not supported')
            else:
                self.stopOtherModel()

    def configModel(self, thread):
        thread.source = self.inputPath
        # 更换模型下 inputPath为None 不能安装模型
        if self.inputPath is not None:
            # 开始按钮状态为运行中|| True
            if self.run_button.isChecked():
                # 首次运行is_continue=True
                if thread.is_continue:
                    thread.start()
                else:
                    thread.resume_thread()
                    self.showStatus('Continue Detection')
            else:
                thread.pause_thread()
                self.showStatus('Pause Detection')
        else:
            self.showStatus("Please select the Image/Video before starting detection...")
            self.run_button.setChecked(False)

    # 运行模型 两个一块运行，一块暂停
    def runModel(self, runbuttonStatus=None):
        self.save_status_button.setEnabled(False)
        if runbuttonStatus:
            self.run_button.setChecked(False)
        # 首先判断是否两边的模型均为正确模型
        if self.checkModelName(self.model_name1) and self.checkModelName(self.model_name2):
            for key in self.left_Threads:
                if self.model_name1 == key:
                    self.configModel(self.left_Threads[key])

            for key in self.right_Threads:
                if self.model_name2 == key:
                    self.configModel(self.right_Threads[key])

        else:
            self.showStatus('The current model is not supported')
            if self.run_button.isChecked():
                self.run_button.setChecked(False)

    # 开始/暂停 预测
    def runorContinue(self):
        if self.inputPath is not None:
            glo.set_value('model_name1', self.model_name1)
            glo.set_value('model_name2', self.model_name2)
            self.reloadModel()
            self.runModel()
        else:
            self.showStatus("Please select the Image/Video before starting detection...")
            self.run_button.setChecked(False)


# 多套一个类 为了实现MouseLabel方法
class VsWindow(VsShow):
    closed = Signal()

    def __init__(self):
        super(VsWindow, self).__init__()
        self.center()
        # --- 拖动窗口 改变窗口大小 --- #
        self.left_grip = CustomGrip(self, Qt.LeftEdge, True)
        self.right_grip = CustomGrip(self, Qt.RightEdge, True)
        self.top_grip = CustomGrip(self, Qt.TopEdge, True)
        self.bottom_grip = CustomGrip(self, Qt.BottomEdge, True)
        # --- 拖动窗口 改变窗口大小 --- #
        self.animation_window = None

    def mousePressEvent(self, event: QMouseEvent) -> None:
        if event.button() == Qt.LeftButton:
            self.mouse_start_pt = event.globalPosition().toPoint()
            self.window_pos = self.frameGeometry().topLeft()
            self.drag = True

    def mouseMoveEvent(self, event: QMouseEvent) -> None:
        if self.drag:
            distance = event.globalPosition().toPoint() - self.mouse_start_pt
            self.move(self.window_pos + distance)

    def mouseReleaseEvent(self, event: QMouseEvent) -> None:
        if event.button() == Qt.LeftButton:
            self.drag = False

    def center(self):
        # PyQt6获取屏幕参数
        screen = QGuiApplication.primaryScreen().size()
        size = self.geometry()
        self.move((screen.width() - size.width()) / 2,
                  (screen.height() - size.height()) / 2 - 10)

    # 拖动窗口 改变窗口大小
    def resizeEvent(self, event):
        # Update Size Grips
        self.resizeGrip()

    def showEvent(self, event):
        super().showEvent(event)
        if not event.spontaneous():
            # 这里定义显示动画
            self.animation = QPropertyAnimation(self, b"windowOpacity")
            self.animation.setDuration(500)  # 动画时间500毫秒
            self.animation.setStartValue(0)  # 从完全透明开始
            self.animation.setEndValue(1)  # 到完全不透明结束
            self.animation.start()

    def closeEvent(self, event):
        if not self.animation_window:
            config_file = 'cfg/setting.json'
            config = dict()
            config['iou'] = self.iou_spinbox.value()
            config['conf'] = self.conf_spinbox.value()
            config['delay'] = self.speed_spinbox.value()
            config['line_thickness'] = self.line_spinbox.value()
            config_json = json.dumps(config, ensure_ascii=False, indent=2)
            with open(config_file, 'w', encoding='utf-8') as f:
                f.write(config_json)
            self.animation_window = QPropertyAnimation(self, b"windowOpacity")
            self.animation_window.setStartValue(1)
            self.animation_window.setEndValue(0)
            self.animation_window.setDuration(500)
            self.animation_window.start()
            self.animation_window.finished.connect(self.close)
            event.ignore()
        else:
            self.setWindowOpacity(1.0)
            self.closed.emit()
